import pandas as pd
import json
import numpy as np
from lifelines import CoxPHFitter
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split


def predict(data):
    data = data.drop('label', axis=1)
    data = data.drop('event', axis=1)
    data = np.array(data)

    file = open('result/params.json', 'r')
    params = json.load(file)
    beta = np.array(params['beta'])
    bch = np.array(params['bch'])
    avg = np.array(params['avg'])
    ph = np.exp(np.sum((data - avg) * beta, axis=1))
    ph = np.expand_dims(ph, axis=0)
    bch = np.expand_dims(bch, axis=1)
    cox = np.exp(-bch * ph)
    survival_result = pd.DataFrame(cox)
    pq_result = survival_result.apply(lambda x: 1-x)
    return pq_result.T


if __name__ == '__main__':
    data_path = '../data/dataset-bia-day-i-cut.csv'
    origin_data = pd.read_csv(data_path)

    # features = pd.read_csv('km/good_feature.csv')
    # features = features.values.tolist()
    # features = [i[0] for i in features]
    # features.extend(['label', 'event'])
    # origin_data = origin_data[features]

    # train_data, test_data = train_test_split(origin_data, test_size=0.2)

    train_data = origin_data[:340]
    # train_data = origin_data
    test_data = origin_data[340:]
    # cox
    cph = CoxPHFitter(penalizer=0.1)
    cph.fit(train_data, duration_col='label', event_col='event', show_progress=True, step_size=0.5)
    cph.print_summary()
    fig, ax = plt.subplots(figsize=(12, 16))
    cph.plot(ax=ax)
    plt.savefig('result/cox.png')
    # cph.check_assumptions(train_data)

    # test
    sf = cph.predict_survival_function(test_data)
    median = cph.predict_median(test_data)
    # ph = cph.predict_partial_hazard(test_data)
    # print(sf)
    # print(median)
    median.to_csv('result/median.csv')
    # print(ph)

    # manual test
    model = dict()

    beta = pd.DataFrame(cph.params_)
    # interval = pd.DataFrame(cph.confidence_intervals_)
    # hr = pd.concat([beta, interval], axis=1)
    # hr.to_csv('weights.csv')
    bch = pd.DataFrame(cph.baseline_cumulative_hazard_)
    avg = test_data.mean()
    avg = avg.drop(['event', 'label'])
    avg = pd.DataFrame(avg, columns=['0'])

    model['beta'] = beta.coef.to_list()
    model['avg'] = avg['0'].to_list()
    model['bch'] = bch['baseline cumulative hazard'].to_list()
    with open('result/params.json', 'w') as f:
        json.dump(model, f)
    #
    result = predict(test_data)
    test_label = test_data.reset_index()
    result_label = pd.concat([result, test_label['label']], axis=1)
    # columns = list(range(1, 41))
    # columns.remove(32)
    # columns.append('label')
    # result_label.columns = columns
    result_label.to_csv('result/result.csv')
